package nlp;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;

import org.tartarus.snowball.ext.porterStemmer;

import weka.core.tokenizers.NGramTokenizer;

public class TfIafFeatureCreator implements FeatureCreator {

	private String textString;
	private porterStemmer stemmer;
	private HashMap<String, Double> tfiafMap;

	private HashSet<String> GetAuthorSet() {
		HashSet<String> authorSet = new HashSet<String>();
		for (String tfiafKey : this.tfiafMap.keySet()) {
			String[] tokens = tfiafKey.split("");
			authorSet.add(tokens[0]);
		}
		return authorSet;
	}

	public TfIafFeatureCreator(String textString, porterStemmer stemmer,
			HashMap<String, Double> tfiafMap) {
		this.textString = textString;
		this.stemmer = stemmer;
		this.tfiafMap = tfiafMap;
	}

	@SuppressWarnings("unchecked")
	@Override
	public void CreateFeature(Object obj) {
		// Tokenize the text.
		int ngramSize = 3;
		NGramTokenizer ngramTokenizer = new NGramTokenizer();
		ngramTokenizer.tokenize(textString);
		ngramTokenizer.setNGramMaxSize(ngramSize);
		ngramTokenizer.setNGramMinSize(ngramSize);

		List<String> ngrams = new ArrayList<String>();
		while (ngramTokenizer.hasMoreElements()) {
			String ngram = (String) ngramTokenizer.nextElement();
			this.stemmer.setCurrent(ngram.toLowerCase());
			this.stemmer.stem();
			String stemmedNgram = stemmer.getCurrent();
			ngrams.add(stemmedNgram);
		}

		HashMap<String, Double> tfiafFeatureMap = new HashMap<String, Double>();
		for (String tfiafKey : this.tfiafMap.keySet()) {
			tfiafFeatureMap.put(tfiafKey, 0.0);
		}

		FreqDistUtil freqDistUtil = new FreqDistUtil(ngrams, ngrams.size());
		List<FreqData> freqDist = freqDistUtil.getAll();
		HashSet<String> authors = GetAuthorSet();
		for (FreqData data : freqDist) {
			for (String author : authors) {
				String authorNgram = author + "" + data.getData();
				if (!this.tfiafMap.containsKey(authorNgram)) {
					continue;
				}
				tfiafFeatureMap.put(authorNgram, data.getFrequency()
						* this.tfiafMap.get(authorNgram));
			}
		}

		HashMap<String, Feature> featureMap = (HashMap<String, Feature>) obj;
		for (String tfiafKey : tfiafFeatureMap.keySet()) {
			Feature feature = new TfIafFeature(tfiafKey, tfiafFeatureMap
					.get(tfiafKey));
			featureMap.put(feature.GetFeatureName(), feature);
		}
	}

	@Override
	public String GetFeatureName() {
		return "TfIaf";
	}

}
